# coding: UTF-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class ModelConfig(object):

    """配置参数"""

    def __init__(self,  num_classes =2 ):
        self.model_name = 'TextCNN'
        self.dropout = 0.5                                              # 随机失活
        self.num_classes = num_classes                        # 类别数
        # 词表大小，在运行时赋值
        self.n_vocab = 400002
        self.embed = 300           # 字向量维度
        self.filter_sizes = (2, 3, 4)                                   # 卷积核尺寸
        self.num_filters = 256   


'''Convolutional Neural Networks for Sentence Classification'''

class TextCNN(nn.Module):
    def __init__(self, embedding = None, freeze_embedding = False, config = ModelConfig(), dim = 768):
        super(TextCNN, self).__init__()
        self.feature_dim = dim
        if embedding is not None:
            self.embedding = nn.Embedding.from_pretrained(
                embedding, freeze=freeze_embedding)
        else:
            self.embedding = nn.Embedding(
                config.n_vocab, config.embed)
        self.convs = nn.ModuleList(
            [nn.Conv2d(1, config.num_filters, (k, config.embed)) for k in config.filter_sizes])
        
    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(self, x):
        out = self.embedding(x)
        out = out.unsqueeze(1)
        feature = out = torch.cat([self.conv_and_pool(out, conv)
                         for conv in self.convs], 1)
        return feature

# def __test__():
#     model = TextCNN()
